package com.moyz.adi.common.rag;

import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.output.Response;

import java.util.List;
import java.util.function.Consumer;

/**
 * AdiChatLanguageModelImpl 是 ChatLanguageModel 的具体实现。
 * <p>
 * 该类组合了一个具体的 ChatLanguageModel 实例，并添加了额外的功能，比如在每次生成响应后执行自定义逻辑。
 */
public class AdiChatLanguageModelImpl implements ChatLanguageModel {

    /**
     * 被包装的具体 ChatLanguageModel 实例。
     */
    private final ChatLanguageModel chatLanguageModel;

    /**
     * 消费者接口，用于处理模型响应。每当有新的 Response<AiMessage> 产生时，
     * 将会调用此消费者来执行额外的操作，如日志记录或数据持久化。
     */
    private final Consumer<Response<AiMessage>> consumer;

    /**
     * 构造函数，初始化 AdiChatLanguageModelImpl 实例。
     *
     * @param chatLanguageModel 具体的 ChatLanguageModel 实例，用于实际的聊天消息处理
     * @param consumer          消费者接口，用于处理来自 chatLanguageModel 的响应
     */
    public AdiChatLanguageModelImpl(ChatLanguageModel chatLanguageModel, Consumer<Response<AiMessage>> consumer) {
        this.chatLanguageModel = chatLanguageModel;
        this.consumer = consumer;
    }

    /**
     * 根据用户的消息生成回复文本。
     *
     * @param userMessage 用户输入的消息文本
     * @return 由语言模型生成的回复文本
     */
    @Override
    public String generate(String userMessage) {
        return chatLanguageModel.generate(userMessage);
    }

    /**
     * 根据一组聊天消息生成 AI 回复，并使用消费者处理响应。
     *
     * @param messages 可变参数列表，包含多个 ChatMessage 对象
     * @return 包含 AIMessage 的响应对象
     */
    @Override
    public Response<AiMessage> generate(ChatMessage... messages) {
        Response<AiMessage> response = chatLanguageModel.generate(messages);
        if (consumer != null) {
            consumer.accept(response); // 使用消费者处理响应
        }
        return response;
    }

    /**
     * 根据聊天消息列表生成 AI 回复，并使用消费者处理响应。
     *
     * @param messages 聊天消息列表
     * @return 包含 AIMessage 的响应对象
     */
    @Override
    public Response<AiMessage> generate(List<ChatMessage> messages) {
        Response<AiMessage> response = chatLanguageModel.generate(messages);
        if (consumer != null) {
            consumer.accept(response); // 使用消费者处理响应
        }
        return response;
    }

    /**
     * 根据聊天消息列表和工具规范列表生成 AI 回复，并使用消费者处理响应。
     *
     * @param messages             聊天消息列表
     * @param toolSpecifications   工具规范列表
     * @return 包含 AIMessage 的响应对象
     */
    @Override
    public Response<AiMessage> generate(List<ChatMessage> messages, List<ToolSpecification> toolSpecifications) {
        Response<AiMessage> response = chatLanguageModel.generate(messages, toolSpecifications);
        if (consumer != null) {
            consumer.accept(response); // 使用消费者处理响应
        }
        return response;
    }

    /**
     * 根据聊天消息列表和单个工具规范生成 AI 回复，并使用消费者处理响应。
     *
     * @param messages            聊天消息列表
     * @param toolSpecification   单个工具规范
     * @return 包含 AIMessage 的响应对象
     */
    @Override
    public Response<AiMessage> generate(List<ChatMessage> messages, ToolSpecification toolSpecification) {
        Response<AiMessage> response = chatLanguageModel.generate(messages, toolSpecification);
        if (consumer != null) {
            consumer.accept(response); // 使用消费者处理响应
        }
        return response;
    }
}
